package org.hong.monkey.rpc

import java.util.concurrent.TimeoutException

import org.hong.monkey.MonkeyConf
import org.hong.monkey.util.Utils

import scala.concurrent.{Await, Awaitable}
import scala.concurrent.duration._

/**
 */
private[rpc] class RpcTimeoutException(message:String, cause:TimeoutException)
  extends TimeoutException(message) { initCause(cause) }

private[monkey] class RpcTimeout(val duration:FiniteDuration, val timeoutProp: String)
  extends Serializable {

  private def createRpcTimeoutException(te: TimeoutException): RpcTimeoutException = {
    new RpcTimeoutException(te.getMessage + ". This timeout is controlled by " + timeoutProp, te)
  }

  def addMessageIfTimeout[T]: PartialFunction[Throwable, T] = {
    case rte: RpcTimeoutException => throw rte
    case te: TimeoutException => throw createRpcTimeoutException(te)
  }

  def awaitResult[T](awaitable: Awaitable[T]): T = {
    try {
      Await.result(awaitable, duration)
    } catch addMessageIfTimeout
  }
}

private[monkey] object RpcTimeout {
  def apply(conf: MonkeyConf, timeoutProp: String): RpcTimeout = {
    val timeout = { conf.getTimeAsSeconds(timeoutProp).seconds }
    new RpcTimeout(timeout, timeoutProp)
  }

  def apply(conf: MonkeyConf, timeoutProp: String, defaultValue:String): RpcTimeout = {
    val timeout = { conf.getTimeAsSeconds(timeoutProp, defaultValue).seconds }
    new RpcTimeout(timeout, timeoutProp)
  }

  def apply(conf: MonkeyConf, timeoutPropList:Seq[String], defaultValue:String): RpcTimeout = {
    require(timeoutPropList.nonEmpty)

    val itr = timeoutPropList.iterator
    var foundProp: Option[(String, String)] = None
    while (itr.hasNext && foundProp.isEmpty) {
      val propKey = itr.next()
      conf.getOption(propKey).foreach { prop => foundProp = Some(propKey, prop) }
    }
    val finalProp = foundProp.getOrElse(timeoutPropList.head, defaultValue)
    val timeout = { Utils.timeStringAsSeconds(finalProp._2).second }
    new RpcTimeout(timeout, finalProp._1)
  }
}
